# -*- coding: utf-8 -*- #

# -----------------------------------------------------------------------
# File Name:    train.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了模型的训练流程
# -----------------------------------------------------------------------

import torch
from torch import nn
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import StepLR
from dataset import CustomDataset
from model import CustomNet


def train_loop(epochs, train_dataloader, val_dataloader, model, loss_fn, optimizer, device):
    """定义训练流程。
    :param epochs: 训练的总轮次
    :param train_dataloader: 训练数据加载器
    :param val_dataloader: 验证数据加载器
    :param model: 模型
    :param loss_fn: 损失函数
    :param optimizer: 优化器
    :param device: 训练设备
    """
    best_accuracy = 0.0
    scheduler = StepLR(optimizer, step_size=10, gamma=0.1)  # 学习率调度器

    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        for batch in train_dataloader:
            images = batch['image'].to(device)
            labels = batch['label'].to(device)

            # 前向传播
            outputs = model(images)
            loss = loss_fn(outputs, labels)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item()

        # 验证阶段
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for batch in val_dataloader:
                images = batch['image'].to(device)
                labels = batch['label'].to(device)
                outputs = model(images)
                val_loss += loss_fn(outputs, labels).item()

                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        # 计算指标
        train_loss /= len(train_dataloader)
        val_loss /= len(val_dataloader)
        accuracy = 100 * correct / total

        # 打印训练信息
        print(f'Epoch [{epoch + 1}/{epochs}], '
              f'Train Loss: {train_loss:.4f}, '
              f'Val Loss: {val_loss:.4f}, '
              f'Val Acc: {accuracy:.2f}%')

        # 保存最佳模型
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), './models/best_model.pth')
            print(f'New best model saved with accuracy: {accuracy:.2f}%')

        # 更新学习率
        scheduler.step()

    # 训练完成后保存最终模型
    torch.save(model.state_dict(), './models/final_model.pth')


if __name__ == "__main__":
    # 定义模型超参数
    BATCH_SIZE = 32
    LEARNING_RATE = 1e-3
    EPOCHS = 30  # 增加到50个epoch

    # 数据预处理
    transform = Compose([
        Resize((64, 64)),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 模型实例化
    model = CustomNet()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 数据加载器
    train_dataloader = DataLoader(
        CustomDataset('./images/train.txt', './images/train', transform),
        batch_size=BATCH_SIZE, shuffle=True)

    val_dataloader = DataLoader(
        CustomDataset('./images/test.txt', './images/test', transform),
        batch_size=BATCH_SIZE)

    # 损失函数和优化器
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)  # 改用Adam优化器

    # 调用训练方法
    train_loop(EPOCHS, train_dataloader, val_dataloader,
               model, loss_fn, optimizer, device)